from zipline.api import symbol, order_target_percent, record, get_datetime
from zipline import run_algorithm
import pandas as pd
import numpy as np

def initialize(context):
    # 策略参数
    context.lookback = 30
    context.zscore_threshold = 2.0
    context.asset_A = symbol('KO')  # 示例资产A
    context.asset_B = symbol('PEP')  # 示例资产B
    
def handle_data(context, data):
    # 获取历史价格数据
    prices_A = data.history(context.asset_A, 'price', context.lookback, '1d')
    prices_B = data.history(context.asset_B, 'price', context.lookback, '1d')
    
    # 计算价差
    spread = prices_A - prices_B
    
    # 计算Z-Score
    mean = np.mean(spread)
    std = np.std(spread)
    zscore = (spread[-1] - mean) / std if std != 0 else 0
    
    # 记录Z-Score
    record(zscore=zscore)
    
    # 交易逻辑
    if zscore > context.zscore_threshold:
        # 做空A，做多B
        order_target_percent(context.asset_A, -0.5)  # 50%仓位
        order_target_percent(context.asset_B, 0.5)   # 50%仓位
    elif zscore < -context.zscore_threshold:
        # 做多A，做空B
        order_target_percent(context.asset_A, 0.5)   # 50%仓位
        order_target_percent(context.asset_B, -0.5)   # 50%仓位
    else:
        # 平仓
        order_target_percent(context.asset_A, 0.0)
        order_target_percent(context.asset_B, 0.0)

# 示例回测配置
if __name__ == '__main__':
    start = pd.Timestamp('2020-01-01', tz='utc')
    end = pd.Timestamp('2023-01-01', tz='utc')
    
    results = run_algorithm(
        start=start,
        end=end,
        initialize=initialize,
        handle_data=handle_data,
        capital_base=10000,
        data_frequency='daily',
        bundle='quandl'
    )